﻿using System;
using System.Collections.Generic;
using Windows.UI.Xaml;
using Windows.UI.Xaml.Media;

namespace Ebt.Mvvm
{
    public static class VisualTreeHelperEx
    {
        public static IEnumerable<DependencyObject> GetAscendants(this DependencyObject element, bool includeItself = false)
        {
            if (element == null)
                throw new ArgumentNullException("element");

            for (var node = includeItself ? element : VisualTreeHelper.GetParent(element);
                 node != null;
                 node = VisualTreeHelper.GetParent(node))
            {
                yield return node;
            }
        }

        public static IEnumerable<DependencyObject> GetDescendants(this DependencyObject d, TraversalKind traversalKind = TraversalKind.BreadthFirst)
        {
            var nodes =
                traversalKind == TraversalKind.BreadthFirst
                ? (ICollection<DependencyObject>)new QueueCollection<DependencyObject>()
                : new StackCollection<DependencyObject>();

            var isInversedOrder = traversalKind != TraversalKind.BreadthFirst;

            nodes.Add(d);

            while (!nodes.IsEmpty)
            {
                var node = nodes.Extract();

                if (!ReferenceEquals(d, node))
                    yield return node;

                foreach (var child in GetChildren(node, isInversedOrder))
                {
                    nodes.Add(child);
                }
            }
        }

        public static IEnumerable<DependencyObject> GetChildren(this DependencyObject d, bool isInversedOrder = false)
        {
            int count = VisualTreeHelper.GetChildrenCount(d);
            if (isInversedOrder)
            {
                for (int i = count - 1; i >= 0; i--)
                {
                    yield return VisualTreeHelper.GetChild(d, i);
                }
            }
            else
            {
                for (int i = 0; i < count; i++)
                {
                    yield return VisualTreeHelper.GetChild(d, i);
                }
            }
        }

        private interface ICollection<TItem>
        {
            bool IsEmpty { get; }
            void Add(TItem item);
            TItem Extract();
            TItem Current { get; }
        }

        private class StackCollection<TItem> : Stack<TItem>, ICollection<TItem>
        {
            #region Implementation of ICollection<out TItem>

            public bool IsEmpty { get { return Count == 0; } }

            public void Add(TItem item)
            {
                Push(item);
            }

            public TItem Extract()
            {
                return Pop();
            }

            public TItem Current
            {
                get { return Peek(); }
            }

            #endregion
        }

        private class QueueCollection<TItem> : Queue<TItem>, ICollection<TItem>
        {
            #region Implementation of ICollection<out TItem>

            public bool IsEmpty { get { return Count == 0; } }

            public void Add(TItem item)
            {
                Enqueue(item);
            }

            public TItem Extract()
            {
                return Dequeue();
            }

            public TItem Current
            {
                get { return Peek(); }
            }

            #endregion
        }
    }

    public enum TraversalKind
    {
        DepthFirstPreOrder,
        //DepthFirstInOrder,
        //DepthFirstPostOrder,
        BreadthFirst,
    }
}
